import numpy as np
import math
import matplotlib.pyplot as plt
from matplotlib import gridspec


def static_vars(**kwargs):
    def decorate(func):
        for k in kwargs:
            setattr(func, k, kwargs[k])
        return func
    return decorate

# -----------------------------------------------
# createRefSignal(y_min, y_max, tau):
#	Function creates the reference signal in
#	range of 10% to 90% of the output range
# -----------------------------------------------
def createRefSignal(y_min, y_max, tau):
	lstRef = []
	width  = int(12*tau)
	nSteps = 8
	arrDy  = [0.1, 0.3, 0.5, 0.7, 0.9, 0.7, 0.5, 0.3]
	rangeY = y_max - y_min
	for idx in arrDy:
		currSP  = idx*rangeY
		lstRef1 = y_min + currSP*np.ones((width))
		lstRef  = np.append(lstRef,lstRef1)
	#lstRef1 = np.sin(np.linspace(0, 38 * np.pi, 12001))
	#lstRef = np.append(lstRef,lstRef1)
	return lstRef

# -----------------------------------------------
# getProcessParameters(processID):
#	Function loads the process parameters 
#	(choosed by operator)
# -----------------------------------------------
def getProcessParameters(processID):
	if processID==1:  # artificial first order process with quadratic static nonlinearity
		U_MIN = 0    					# actuator's lower limit
		U_MAX = 10 						# actuator's upper limit
		Y_MIN = 0    					# colntrolled variable's lower limit
		Y_MAX = 10 						# colntrolled variable's upper limit
		TAU   = 40   					# Time constant in seconds
		T_S   = 1    					# sampling time
		A_R   = 1-T_S/TAU				# Pole of the reference model

	elif processID==2: # a hydraulic pilot plant of two tanks
		U_MIN = 0    					# actuator's lower limit
		U_MAX = 5  						# actuator's upper limit
		Y_MIN = 0    					# colntrolled variable's lower limit
		Y_MAX = 5  						# colntrolled variable's upper limit
		TAU   = 40   					# Time constant in seconds
		T_S   = 1    					# sampling time
		A_R   = 1-T_S/TAU				# Pole of the reference model	

	elif processID==3: # yourself
		U_MIN = -5    					# actuator's lower limit
		U_MAX = 5  						# actuator's upper limit
		Y_MIN = -5    					# colntrolled variable's lower limit
		Y_MAX = 5  						# colntrolled variable's upper limit
		TAU   = 40   					# Time constant in seconds
		T_S   = 1    					# sampling time
		A_R   = 1-T_S/TAU				# Pole of the reference model	

	else:
		print("---------------------------------------------------------------------")
		print("WARRNING: Please choose valind values for the variable 'processID'.")
		print("---------------------------------------------------------------------\n")

	return U_MIN, U_MAX, Y_MIN, Y_MAX, TAU, T_S, A_R
	
	
# -----------------------------------------------
# fnProcessModel(idxModel, u, y, idxK):
#	Function calculates the process output
# -----------------------------------------------
@static_vars(stateVar=0)
def fnProcessModel(idxModel, u, y, idxK):
	if idxModel==1: # artificial first order process with quadratic static nonlinearity
		a_p   = 0.98  
		b_p   = 0.01
		mu    = 0
		sig   = 0.005
		noise = np.random.normal(mu, sig, 1)
		y_p   = fnProcessModel.stateVar
		if idxK == 1:
			y_p = a_p*math.sqrt(abs(y)) + b_p*u
			y   = y_p*y_p + noise
		else:
			y_p = a_p*y_p + b_p*u
			y   = y_p*y_p + noise

		fnProcessModel.stateVar = y_p

	elif idxModel==2: # a hydraulic pilot plant of two tanks
		kv1   = 3;
		kv2   = 1;
		a_p   = 1/50;
		b_p   = 1/50;
		dt    = 2;
		mu    = 0
		sig   = 0.005
		noise = np.random.normal(mu, sig, 1)
		y_p   = fnProcessModel.stateVar
		if idxK == 1:
			y_p = y_p + dt*a_p*(u - kv1*math.sqrt(abs(y_p-y))*np.sign(y_p-y))
			y   = y + dt*b_p*(kv1*math.sqrt(abs(y_p-y))*np.sign(y_p-y) - kv2*math.sqrt(y))
		else:
			y_p = y_p + dt*a_p*(u - kv1*math.sqrt(abs(y_p-y))*np.sign(y_p-y))
			y   = y + dt*b_p*(kv1*math.sqrt(abs(y_p-y))*np.sign(y_p-y) - kv2*math.sqrt(abs(y)))

		y = abs(y) + noise
		fnProcessModel.stateVar = y_p

	elif idxModel==3: # yourself
		# y=pos, y_p=vel
		g = 9.8
		l = 0.5
		m = 0.1
		mc = 1
		mu    = 0
		sig   = 0.005
		T = 0.01
		noise = np.random.normal(mu, sig, 1)
		y_p = fnProcessModel.stateVar

		k11 = y_p
		k21 = (g*math.sin(y) - m*l*math.pow(y_p, 2)*math.cos(y)*math.sin(y)/(mc + m)) / (l*(4/3 - m*math.pow(math.cos(y), 2)/(mc + m))) + (math.cos(y)/(mc + m)) / (l*(4/3 - m*math.pow(math.cos(y), 2)/(mc + m))) * u
		k12 = y_p + k21*T/2
		k22 = (g*math.sin(y + k11*T/2) - m*l*math.pow(y_p + k21*T/2, 2)*math.cos(y + k11*T/2)*math.sin(y + k11*T/2)/(mc + m)) / (l*(4/3 - m*math.pow(math.cos(y + k11*T/2), 2)/(mc + m))) + (math.cos(y + k11*T/2)/(mc + m)) / (l*(4/3 - m*math.pow(math.cos(y + k11*T/2), 2)/(mc + m))) * u
		k13 = y_p + k22*T/2
		k23 = (g*math.sin(y + k12*T/2) - m*l*math.pow(y_p + k22*T/2, 2)*math.cos(y + k12*T/2)*math.sin(y + k12*T/2)/(mc + m)) / (l*(4/3 - m*math.pow(math.cos(y + k12*T/2), 2)/(mc + m))) + (math.cos(y + k12*T/2)/(mc + m)) / (l*(4/3 - m*math.pow(math.cos(y + k12*T/2), 2)/(mc + m))) * u
		k14 = y_p + k23*T
		k24 = (g*math.sin(y + k13*T) - m*l*math.pow(y_p + k23*T, 2)*math.cos(y + k13*T)*math.sin(y + k13*T)/(mc + m)) / (l*(4/3 - m*math.pow(math.cos(y + k13*T), 2)/(mc + m))) + (math.cos(y + k13*T)/(mc + m)) / (l*(4/3 - m*math.pow(math.cos(y + k13*T), 2)/(mc + m))) * u

		y = y + T*(k11 + 2*k12 + 2*k13 + k14)/6
		y_p = y_p + T*(k21 + 2*k22 + 2*k23 + k24)/6
		fnProcessModel.stateVar = y_p

	return y
	
# -----------------------------------------------
# fnPlotResults(lenSim, buff..., myEvolve):
#	Function plots the results
# -----------------------------------------------
def fnPlotResults(lenSim, idxK, buff_Y, buff_U, buff_Y_R, buff_R, buff_EPS, buff_Data, buff_Theta, myEvolve):
	fig = plt.figure(figsize=(14,10))
	gs  = gridspec.GridSpec(9, 8)
	plt.subplot(gs[0:2,0:2])
	plt.plot(buff_R[1:lenSim],   label="Reference")
	plt.plot(buff_Y_R[1:lenSim], label="Model reference")
	plt.plot(buff_Y[1:lenSim],   label="Controlled singal")
	plt.gca().axes.xaxis.set_ticklabels([])
	plt.title('STARTING PHASE')
	plt.legend()

	plt.subplot(gs[0:2,2:4])
	plt.plot(buff_R[(idxK-lenSim+1):idxK],   label="Reference")
	plt.plot(buff_Y_R[(idxK-lenSim+1):idxK], label="Model reference")
	plt.plot(buff_Y[(idxK-lenSim+1):idxK],   label="Controlled singal")
	plt.gca().axes.xaxis.set_ticklabels([])
	plt.gca().axes.yaxis.set_ticklabels([])
	plt.title('FINISHING PHASE')
	plt.legend()

	plt.subplot(gs[2:4,0:2])
	plt.plot(buff_U[1:lenSim])
	plt.gca().axes.xaxis.set_ticklabels([])
	plt.ylabel('Control signal')

	plt.subplot(gs[2:4,2:4])
	plt.plot(buff_U[(idxK-lenSim+1):idxK])
	plt.gca().axes.xaxis.set_ticklabels([])
	plt.gca().axes.yaxis.set_ticklabels([])

	plt.subplot(gs[4,:])
	plt.plot(buff_EPS, '-y')
	plt.gca().axes.xaxis.set_ticklabels([])
	plt.ylabel('Tracking error')

	plt.subplot(gs[0:4,4:8])
	for ii in range(0, len(myEvolve.cloudList)):
		hPlot = plt.plot(buff_Data[myEvolve.cloudList[ii].T_k,0], buff_Data[myEvolve.cloudList[ii].T_k,1], '.')
		plt.plot(myEvolve.cloudList[ii].mi[0], myEvolve.cloudList[ii].mi[1], 'o', markersize=20.0, color=hPlot[-1].get_color(), mec='k', mew=3.0)
		plt.gca().axes.yaxis.set_ticks_position('right')
		plt.gca().axes.xaxis.set_ticklabels([])
	plt.title('Clouds')
	
	labelPID = ['P', 'I', 'D', 'R']
	iTheta = -1
	for iCloud in range(0, len(myEvolve.cloudList)):		
		iSubPlot = 4
		for iParam in range(0, 4):
			iSubPlot += 1
			iTheta   += 1
			plt.subplot(gs[iSubPlot,:])
			plt.plot(buff_Theta[:,iTheta])
			plt.ylabel(labelPID[iParam])
			if iParam < 3:
				plt.gca().axes.xaxis.set_ticklabels([])
			#print('iTheta:  ', iTheta)
			#plt.ylabel('Control signal')


	plt.subplots_adjust(left=0.05, bottom=0.05, right=0.95, top=0.95, wspace=0.05, hspace=0.10)
	plt.show()
	
	